/*
    This program is free software: you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation, either version 3 of the License, or
    (at your option) any later version.

    This program is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "NTree.h"
#include <string>
#include <fstream>
#include "AuxMath.h"
#include "Vector2i.h"

#pragma once

template<class T>
class Bitmap
{
private:
	static const unsigned int wordWidthBytes=sizeof(unsigned int);
	static const unsigned int wordWidthBits=wordWidthBytes<<3;
	struct BitmapHeader
	{
		uint32_t filesize;
		uint16_t res1;
		uint16_t res2;
		uint32_t pixelOffset;
	};

	struct BitmapInfoHeader
	{
		uint32_t size;
		uint32_t width;
		uint32_t height;
		uint16_t colorPlanes;
		uint16_t bpp;
		uint32_t compression;
		uint32_t rawSize;
		uint32_t pxPerMeterX;
		uint32_t pxPerMeterY;
		uint32_t paletteColors;
		uint32_t colors;
	};

	static void
	rasterSubVoxels(const typename NTree<Vector2i, T>::Node& node,
			const Area<Vector2i>& area, unsigned int* puiBits,
			const unsigned int& width, const bool bBlocks=false)
	{
		const Vector2i vShift(width>>1);
		Box<Vector2i> areaBox=area.toBox();
		if(bBlocks)
			areaBox=Box<Vector2i>(areaBox.min.doubleVec(), areaBox.max.doubleVec());

		switch(node.getType())
		{
		case Nil:
			break;
		case Solid:
			//write rectangle to bitmap
			if(bBlocks)
				setBits(puiBits, width, width,
					Box<Vector2i>(areaBox.min+vShift, areaBox.max+vShift/*-Vector2i(1)*/));
			else
				setBits(puiBits, width, width,
					Box<Vector2i>(areaBox.min+vShift, areaBox.max+vShift));
			break;
		default:
			Area<Vector2i> subAreas[Vector2i::childCount];
			area.subAreas(subAreas);
			for (unsigned int i = 0; i < Vector2i::childCount; i++)
				rasterSubVoxels(node.constGetSubNode(i), subAreas[i],
						puiBits, width, bBlocks);
			break;
		}
	}

	static __attribute__((const))
	unsigned int
	getRowSize(const unsigned int width)
	{
		const unsigned int nRemainderBits=width&(wordWidthBits-1);
		unsigned int uiRegularRowSize=width/wordWidthBits;
		unsigned int uiRowSize=uiRegularRowSize;
		if(nRemainderBits)
			uiRowSize++;
		return uiRowSize;
	}

public:
	static unsigned int*
	allocBits(const unsigned int width, const unsigned int height)
	{
		const unsigned int uiRowSize=getRowSize(width);

		unsigned int* pBits=new unsigned int[height * uiRowSize];
		memset(pBits, 0, sizeof(unsigned int) * height * uiRowSize);
		return pBits;
	}

	static void
	freeBits(unsigned int* p)
	{
		delete[] p;
	}

	static int
	treeBase(const NTree<Vector2i, T>& tree)
	{
		return AuxMath::pow2(tree.getSize())<<1;
	}

	static void
	setBits(unsigned int* puiBits,
			const unsigned int& width,
			const unsigned int& height,
			const Box<Vector2i>& rect)
	{
		const unsigned int uiRowSize=getRowSize(width);

		for (int y = rect.min.y; y <= rect.max.y; y++)
		{
			for (int x = rect.min.x; x <= rect.max.x; x++)
			{
				const unsigned int uiWord=x/wordWidthBits;
				const unsigned int uiBit=x & (wordWidthBits-1);

				puiBits[y * uiRowSize + uiWord] |= 1<<(wordWidthBits-1 - uiBit);
			}
		}
	}

	static bool
	getBit(const unsigned int* puiBits,
			const unsigned int& width,
			const unsigned int& height,
			const Vector2i v)
	{
		if(puiBits==nullptr)
			return false;
		const unsigned int uiRowSize=getRowSize(width);
		const unsigned int uiWord=v.x/wordWidthBits;
		const unsigned int uiBit=v.x & (wordWidthBits-1);

		return puiBits[v.y * uiRowSize + uiWord] & (1<<(wordWidthBits-1 - uiBit));
	}

	static void
	treeToBits(const NTree<Vector2i, T>& tree, unsigned int* puiBits,
			const bool bBlocks=false)
	{
		unsigned int uiBase=treeBase(tree);
		if(bBlocks)
			uiBase<<=1;
		Area<Vector2i> subs[Vector2i::childCount];
		tree.subAreas(subs);

		for (unsigned int i = 0; i < Vector2i::childCount; i++)
			rasterSubVoxels(tree.constGetRoot(i), subs[i],
				puiBits, uiBase, bBlocks);
	}

	static bool
	exportBmp(unsigned int* puiBits, unsigned int width, unsigned int height, std::string filename)
	{
		//check file extension
		static const std::string strExt(".bmp");
		if(filename.size()<=4)
		//too short for extension
			filename += strExt;
		else
		{
			if(filename.compare(filename.size()-strExt.size(), strExt.size(), strExt))
				//different ending
				filename += strExt;
		}
		BitmapHeader bmp;
		BitmapInfoHeader bmpInfo;
		static const unsigned int bpp=1;
		static const unsigned int colors=1<<bpp;

		bmpInfo.size=sizeof(bmpInfo);
		bmpInfo.width=width;
		bmpInfo.height=height;
		bmpInfo.colorPlanes=1;
		bmpInfo.bpp=bpp;
		bmpInfo.compression=0;
		bmpInfo.pxPerMeterX=bmpInfo.pxPerMeterY=4000;
		bmpInfo.paletteColors=colors;
		bmpInfo.colors=colors;

		unsigned int uiStride=(((bmpInfo.width-1)>>5) + 1)<<2; //in bytes; aligns to 32bit
		unsigned int uiRawSize=uiStride*bmpInfo.height;

		char aID[2]={'B', 'M'};
		bmp.pixelOffset=sizeof(bmpInfo)+sizeof(bmp)+sizeof(aID)+bmpInfo.paletteColors*4;
		bmp.filesize=bmp.pixelOffset+uiRawSize;

		std::ofstream out(filename);

		out.write(aID, sizeof(aID));
		out.write((char*)&bmp, sizeof(bmp));
		out.write((char*)&bmpInfo, sizeof(bmpInfo));
		//write palette
		static const uint32_t blk=0,
				grn=0x00cccccc;
		out.write((const char*)&blk, 4);
		out.write((const char*)&grn, 4);

		//stride of input in integers
		const unsigned int nRemainderBits=width&(wordWidthBits-1);
		unsigned int uiRegularRowSize=width/wordWidthBits;
		unsigned int uiRowSize=uiRegularRowSize;
		if(nRemainderBits)
			uiRowSize++;

		unsigned char* pcRowBuffer=new unsigned char[uiStride];

		for(unsigned int y=0; y<height; y++)
		{
			for(unsigned int x=0; x<uiRegularRowSize; x++)
			{
				const unsigned int uiBits=puiBits[y*uiRowSize+x];
				//copy bytes
				for(unsigned int iByte=0; iByte<wordWidthBytes; iByte++)
				{
					unsigned int uiByte=(uiBits>>((wordWidthBytes-iByte-1)<<3)) & 0xff;
					pcRowBuffer[iByte+(x*wordWidthBytes)]=(unsigned char)uiByte;
				}
			}
			if(nRemainderBits)
			{
				const unsigned int uiBits=puiBits[y*uiRowSize+uiRegularRowSize];
				const unsigned int offsetByte=uiRegularRowSize*wordWidthBytes;
				unsigned int remainingBytes=uiStride-offsetByte;

				for(unsigned int iByte=0; iByte<remainingBytes; iByte++)
				{
					unsigned int uiByte=(uiBits>>((wordWidthBytes-iByte-1)<<3)) & 0xff;
					pcRowBuffer[offsetByte+iByte]=uiByte;
				}
			}
			out.write((char*)pcRowBuffer, uiStride);
		}

		delete pcRowBuffer;
		out.close();

		return true;
	}

	static bool
	importBmp(std::string filename, unsigned int*& puiBits, unsigned int& width, unsigned int& height)
	{
		//check file extension
		puiBits=nullptr;

		BitmapHeader bmp;
		BitmapInfoHeader bmpInfo;
		char aID[2];

		static const unsigned int bpp=1;

		std::ifstream in(filename);

		in.read(aID, sizeof(aID));
		if(*aID != 'B' || aID[1]!='M')
			return false;

		in.read((char*)&bmp, sizeof(bmp));
		in.read((char*)&bmpInfo, sizeof(bmpInfo));

		width=bmpInfo.width;
		height=bmpInfo.height;
		if(bmpInfo.colorPlanes!=1)
			return false;
		if(bmpInfo.bpp!=bpp)
			return false;
		if(bmpInfo.compression!=0)
			return false;

		if(bmpInfo.size!=sizeof(bmpInfo))
			in.seekg(bmpInfo.size-sizeof(bmpInfo), std::ios_base::cur);

		in.seekg(bmp.pixelOffset, std::ios_base::beg);

		unsigned int uiStride=(((bmpInfo.width-1)>>5) + 1)<<2; //in bytes; aligns to 32bit
		unsigned int uiRawSize=uiStride*bmpInfo.height;

		if(bmp.filesize!=(bmp.pixelOffset+uiRawSize))
			return false;

		//stride of input in integers
		const unsigned int nRemainderBits=width&(wordWidthBits-1);
		unsigned int uiRegularRowSize=width/wordWidthBits;
		unsigned int uiRowSize=uiRegularRowSize;
		if(nRemainderBits)
			uiRowSize++;

		puiBits=allocBits(width, height);

		unsigned char* pcRowBuffer=new unsigned char[uiStride];

		for(unsigned int y=0; y<height; y++)
		{
			in.read((char*)pcRowBuffer, uiStride);

			for(unsigned int x=0; x<uiRegularRowSize; x++)
			{
				unsigned int uiBits=0;
				//copy bytes
				for(unsigned int iByte=0; iByte<wordWidthBytes; iByte++)
				{
					unsigned int uiByte=pcRowBuffer[iByte+(x*wordWidthBytes)];
					uiBits|=uiByte<<((wordWidthBytes-iByte-1)<<3);
				}
				puiBits[y*uiRowSize+x]=uiBits;
			}
			if(nRemainderBits)
			{
				unsigned int uiBits=0;
				const unsigned int offsetByte=uiRegularRowSize*wordWidthBytes;
				unsigned int remainingBytes=uiStride-offsetByte;

				for(unsigned int iByte=0; iByte<remainingBytes; iByte++)
				{
					unsigned int uiByte=pcRowBuffer[offsetByte+iByte];
					uiBits|=uiByte<<((wordWidthBytes-iByte-1)<<3);
				}
				puiBits[y*uiRowSize+uiRegularRowSize]=uiBits;
			}
		}

		delete pcRowBuffer;
		in.close();

		return true;
	}

	static bool
	treeToBitmap(std::string filename, const NTree<Vector2i, T>& tree,
			const bool bBlocks=false)
	{
		unsigned int uiBase=treeBase(tree);
		if(bBlocks)
			uiBase<<=1;
		unsigned int* p=allocBits(uiBase, uiBase);
		treeToBits(tree, p, bBlocks);
		if(!exportBmp(p, uiBase, uiBase, filename))
			return false;
		freeBits(p);
		return true;
	}

	static bool
	writeComplementaryBitmaps(std::string filename, const NTree<Vector2i, T>& tree)
	{
		treeToBitmap(filename, tree, true);
		CopyCon<Vector2i, T> con(tree);
		NTree<Vector2i, T> t(con, T());
		NTreeHelpers::invertVoxels<Vector2i, int>(t);
		treeToBitmap(filename + std::string("_inv"), t, true);
		return true;
	}

	static bool
	bitmapToTree(NTree<Vector2i, T>& tree, std::string filename)
	{
		unsigned int* puiBits;
		unsigned int width, height;
		if(!importBmp(filename, puiBits, width, height))
		{
			printf("error opening %s!\n", filename.c_str());
			return false;
		}

		BitmapCon bmpCon(puiBits, width, height);

		tree=NTree<Vector2i, T>(bmpCon, T());

		freeBits(puiBits);
		return true;
	}

	class BitmapCon
	{
	public:
		BitmapCon(const unsigned int* puiBits,
				const unsigned int& width,
				const unsigned int& height)
		: puiBits(puiBits), width(width), height(height),
		  bound(Vector2i(width-1, height-1)),
		  vShift(bound.half()),
		  aabb(Box<Vector2i>(vShift.reverse(), bound-vShift)),
		  size(AuxMath::log2(aabb.max.maxComponent()))
		{

		}

		const int getSize() const {return size;}

		VoxelType fn(const Box<Vector2i>& box) const
		{
			if(BoxTests<Vector2i>::boxRelation(box, aabb)==Outside)
				return Nil;
			if(box.min!=box.max)
				return Both;

			return (getBit(puiBits, width, height, box.min+vShift) ? Solid : Nil);
		}
	private:
		const unsigned int* puiBits;
		const unsigned int width;
		const unsigned int height;

		const Vector2i bound;
		const Vector2i vShift;
		const Box<Vector2i> aabb;

		const unsigned int size;
	};
};
